import os
import random
from random import shuffle

import cv2
import numpy as np
from PIL import Image
from torch.utils.data.dataset import Dataset


def rand(a=0, b=1):
    return np.random.rand()*(b-a) + a

#继承Dataset
class SiameseDataset(Dataset):
    def __init__(self, input_shape, dataset_path, num_train, num_val, train=True, train_own_data=True):
        super(SiameseDataset, self).__init__()

        self.dataset_path   = dataset_path
        self.image_height   = input_shape[0]
        self.image_width    = input_shape[1]
        self.channel        = input_shape[2]
        
        self.train_lines    = []#顺序存放图片地址
        self.train_labels   = []#按照图片顺序存放的种类标签

        self.val_lines      = []
        self.val_labels     = []
        self.types          = 0#种类数

        self.test_coute = 0

        self.num_train      = num_train
        self.num_val        = num_val
        self.train          = train
        self.train_own_data = train_own_data
        self.load_dataset()

    def __len__(self):
        if self.train:
            # train = self.num_train
            return int(self.num_train)
        else:
            return int(self.num_val)

    def load_dataset(self):
        # train_path = os.path.join(self.dataset_path, 'images_background')
        if self.train_own_data:
            #-------------------------------------------------------------#
            #   Q -> 0
            #   Z -> 1
            #-------------------------------------------------------------#
            # file = os.listdir(self.dataset_path)
            # file.sort(key=lambda x:int(x))
            Q_path = os.path.join(self.dataset_path, 'Q')
            for image in os.listdir(Q_path):
                self.train_lines.append(os.path.join(Q_path, image))
                self.train_labels.append(0)

            Z_path = os.path.join(self.dataset_path, 'Z')
            for image in os.listdir(Z_path):
                self.train_lines.append(os.path.join(Z_path, image))
                self.train_labels.append(1)

            c = np.column_stack((self.train_lines, self.train_labels))  # 将y添加到x的最后一列 4*4
            np.random.shuffle(c)

            self.train_lines = c[:, :-1].flatten()  # 乱序后的x
            self.train_labels = c[:, -1]  # 同等乱序后的y

            # for character in file:
            #     #-------------------------------------------------------------#
            #     #   对每张图片进行遍历
            #     #-------------------------------------------------------------#
            #     character_path = os.path.join(train_path, character)
            #     for image in os.listdir(character_path):
            #         self.train_lines.append(os.path.join(character_path, image))
            #         self.train_labels.append(self.types)
            #     self.types += 1
            #     t = self.train_lines
            #     m = self.train_labels
            # print(self.train_lines)
        
        #-------------------------------------------------------------#
        #   将训练集和验证集进行划分
        #-------------------------------------------------------------#
        self.val_lines      = self.train_lines[self.num_train:]
        self.val_labels     = self.train_labels[self.num_train:]
    
        self.train_lines    = self.train_lines[:self.num_train]
        self.train_labels   = self.train_labels[:self.num_train]

    def get_random_data(self, image, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, flip_signal=False):
        if self.channel == 1:
            image = image.convert("RGB")

        h, w = input_shape
        #------------------------------------------#
        #   图像大小调整
        #------------------------------------------#
        rand_jit1   = rand(1-jitter,1+jitter)
        rand_jit2   = rand(1-jitter,1+jitter)
        new_ar      = w/h * rand_jit1/rand_jit2#变化后图片的长宽比例

        scale = rand(0.75,1.25)#这里就是图片大小在75%到125%之间
        if new_ar < 1:
            nh = int(scale*h)
            nw = int(nh*new_ar)
        else:
            nw = int(scale*w)
            nh = int(nw/new_ar)
        image = image.resize((nw,nh), Image.BICUBIC)#三次样条插值
        
        #------------------------------------------#
        #   翻转图像
        #------------------------------------------#
        flip = rand()<.5#按照概率翻转
        if flip and flip_signal: 
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
        #------------------------------------------#
        #   随机遮盖图像
        #------------------------------------------#
        dx = int(rand(0, w-nw))
        dy = int(rand(0, h-nh))
        new_image = Image.new('RGB', (w,h), (255,255,255))
        new_image.paste(image, (dx, dy))
        image = new_image
        #------------------------------------------#
        #   图像旋转
        #------------------------------------------#
        rotate = rand()<.5
        if rotate: 
            angle = np.random.randint(-5,5)
            a,b = w/2,h/2
            M = cv2.getRotationMatrix2D((a,b),angle,1)
            image = cv2.warpAffine(np.array(image), M, (w,h), borderValue = [255,255,255])
        #------------------------------------------#
        #   色域扭曲 转到HSV色域调整
        #------------------------------------------#
        hue = rand(-hue, hue)
        sat = rand(1, sat) if rand()<.5 else 1/rand(1, sat)
        val = rand(1, val) if rand()<.5 else 1/rand(1, val)
        x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV)
        x[..., 0] += hue*360
        x[..., 0][x[..., 0]>1] -= 1
        x[..., 0][x[..., 0]<0] += 1
        x[..., 1] *= sat
        x[..., 2] *= val
        x[x[:,:, 0]>360, 0] = 360
        x[:, :, 1:][x[:, :, 1:]>1] = 1
        x[x<0] = 0
        image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
        if self.channel == 1:
            image_data = Image.fromarray(np.uint8(image_data)).convert("L")

        return image_data

    #
    def _convert_path_list_to_images_and_labels(self, path):
        #-------------------------------------------#
        #   这里是2，number_of_pairs队数，两队，一正一反
        #-------------------------------------------#
        # number_of_pairs = int(len(path_list) / 2)
        #-------------------------------------------#
        #   定义网络的输入图片和标签 这里是pairs_of_images(2,3,h,w) label(2,1)
        #-------------------------------------------#
        # pairs_of_images = [np.zeros(self.channel, self.image_height, self.image_width) ]
        # labels          = np.zeros((number_of_pairs, 1))

        # -------------------------------------------#
        #   将图片0填充到输入[0][0]中
        # -------------------------------------------#
        image = Image.open(path)
        image = self.get_random_data(image, [self.image_height, self.image_width])
        image = np.asarray(image).astype(np.float64)
        image = np.transpose(image, [2, 0, 1])  # 变换通道
        image = image / 255  # 将值置0，1
        # pairs_of_images[ :, :, :] = image
        # pairs_of_images[0][0, :, :, :] = image



        return image

    #提供索引
    def __getitem__(self, index):
        if self.train:
            #lines存放path,label存放种类
            lines = self.train_lines
            labels = self.train_labels
        else:
            lines = self.val_lines
            labels = self.val_labels

        batch_images_path = lines[index]
        # print("sel-",batch_images_path)
        labels = int(labels[index])
        #
        images = self._convert_path_list_to_images_and_labels(batch_images_path)
        #这里的imgage(2,2,3,105,105),其中的【0，2,3，105，105]存放的是相同的两张图片，反之为不同的图片
        return images, labels

# DataLoader中collate_fn使用
# 这里的batch是存储一个batch_size的图像（一组包含2张图）
# 1表示相同，0表示不同
# def dataset_collate(batch):
#     left_images     = []
#     right_images    = []
#     labels          = []
#
#     #pair_imgs包含一组图片
#     for pair_imgs, pair_labels in batch:
#         for i in range(len(pair_imgs[0])):
#             # m = pair_imgs[0][i]#两张图片
#             # n = pair_imgs[1][i]
#             left_images.append(pair_imgs[0][i])
#             right_images.append(pair_imgs[1][i])#这里存放了几组不同的
#             labels.append(pair_labels[i])
#     # print("dataset_collate")
#     return np.array([left_images, right_images]), np.array(labels)
